package com.ubxtech.gateway.filter;

import com.ubxtech.gateway.balance.GrayLoadBalancerRule;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.gateway.config.GatewayLoadBalancerProperties;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.ReactiveLoadBalancerClientFilter;
import org.springframework.cloud.gateway.support.DelegatingServiceInstance;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.net.URI;

/**
 * 灰度过滤器,自定义负载均衡的过滤器
 * @author Victor.Xiao
 * @since 2023-08-22 14:46
*/
@Slf4j
public class GrayReactiveLoadBalancerFilter extends ReactiveLoadBalancerClientFilter {
	private GatewayLoadBalancerProperties properties;

	private GrayLoadBalancerRule grayLoadBalancerRule;

	public GrayReactiveLoadBalancerFilter(GatewayLoadBalancerProperties properties, GrayLoadBalancerRule grayLoadBalancerRule) {
		super(null, null);
		this.properties = properties;
		this.grayLoadBalancerRule = grayLoadBalancerRule;
	}

	@Override
	public int getOrder() {
		return LOAD_BALANCER_CLIENT_FILTER_ORDER;
	}

	@Override
	public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
		//=====基本是复制ReactiveLoadBalancerClientFilter中的逻辑====
		URI url = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
		String schemePrefix = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR);
		//判断url 前缀 如不是lb开头的就进行下一个过滤器
		if (url == null || (!"lb".equals(url.getScheme()) && !"lb".equals(schemePrefix))) {
			return chain.filter(exchange);
		}
		//根据网关的原始网址。替换exchange url为 http://IP:PORT/path 路径的url
		// preserve the original url
		ServerWebExchangeUtils.addOriginalRequestUrl(exchange, url);

		if (log.isDebugEnabled()) {
			log.debug(" url before > GrayReactiveLoadBalancerFilter: " + url);
		}

		return choose(exchange).doOnNext(response -> {
			//=====基本是复制ReactiveLoadBalancerClientFilter中的逻辑====
			if (!response.hasServer()) {
				throw NotFoundException.create(properties.isUse404(), "Unable to find instance for > GrayReactiveLoadBalancerFilter: " + url.getHost());
			}

			URI uri = exchange.getRequest().getURI();

			// if the `lb:<scheme>` mechanism was used, use `<scheme>` as the default,
			// if the loadbalancer doesn't provide one.
			ServiceInstance retrievedInstance = (ServiceInstance)response.getServer();
			String overrideScheme = retrievedInstance.isSecure() ? "https" : "http";
			if (schemePrefix != null) {
				overrideScheme = url.getScheme();
			}

			DelegatingServiceInstance serviceInstance = new DelegatingServiceInstance(
					response.getServer(), overrideScheme);

			URI requestUrl = LoadBalancerUriTools.reconstructURI(serviceInstance, uri);

			if (log.isDebugEnabled()) {
				log.debug("LoadBalancerClientFilter url chosen  > GrayReactiveLoadBalancerFilter:" + requestUrl);
			}
			exchange.getAttributes().put(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR, requestUrl);
		}).then(chain.filter(exchange));
	}

	/**
	 * 选择自定义的负载均衡
	 * @author Victor.Xiao
	 * @since 2023-08-22 15:19
	*/
	private Mono<Response<ServiceInstance>> choose(ServerWebExchange exchange) {
		URI uri = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
		ServiceInstance serviceInstance = grayLoadBalancerRule.choose(uri.getHost(), exchange.getRequest());
		return Mono.just(new DefaultResponse(serviceInstance));
	}
}
